import numpy as np
from PIL import Image
import glob
import argparse

def parse_args():
    parser = argparse.ArgumentParser(
        description='if calc dir, the img1 and img2 must have the same name')
    parser.add_argument('--image', default=False)
    parser.add_argument('--img1_dir', default='')
    parser.add_argument('--img2_dir', default='')
    args = parser.parse_args()
    return args

width = 1280
height = 768

# 读取图片并转换为array
def getImage(path):
    img = Image.open(path)
    # img = img.resize((1024, 608), Image.Resampling.BILINEAR)
    img = img.resize((height, width), Image.BILINEAR)
    return np.array(img)


# 计算图像array的差值，并返回list
def getDiff(image_list):
    res_mean = []
    res_var = []
    for i in range(1, len(image_list)):
        diff = np.abs(image_list[i] - image_list[i - 1]) / np.sqrt(608 * 1024 * 3)
        res_mean.append(np.mean(diff))
        res_var.append(np.var(diff))
    return res_mean, res_var


# 计算DiFF的均值和方差
def get_mean_var(diff_mean, diff_var):
    mean = np.mean(diff_mean)
    var = np.mean(diff_var)
    return mean, var


if __name__ == '__main__':
    args = parse_args()

    root_gt = args.img1_dir
    root = args.img2_dir
    image_path_list_gt = sorted(glob.glob(root_gt + '/*.*'))
    print("# of image pairs:", len(image_path_list_gt))
    image_array_list = []
    for image_path in image_path_list_gt:
        image_array_list.append(getImage(image_path))
    diff_mean, diff_var = getDiff(image_array_list)
    mean_gt, var_gt = get_mean_var(diff_mean, diff_var)
    print("dir1 mean = %3.4f , var = %3.4f" % (mean_gt, var_gt))

    image_path_list = sorted(glob.glob(root + '/*.*'))[1:]
    image_array_list = []
    for image_path in image_path_list:
        image_array_list.append(getImage(image_path))
    diff_mean, diff_var = getDiff(image_array_list)
    mean, var = get_mean_var(diff_mean, diff_var)
    print("dir2 mean = %3.4f var = %3.4f" % (mean, var))
    print("diff mean = %3.4f diff var = %3.4f " % (np.abs(mean-mean_gt), np.abs(var-var_gt)))
    print('----------------------------------------')
